import time

import pyautogui
import cv2
import numpy as np
from tkinter import Tk, Button

import pickle
import os


class ScreenshotGetPosition:
    def __init__(self, point_number):
        self.point_number = point_number

        self.folder = "setting"
        self.file_name = "data.pkl"
        self.data_path = os.path.join(self.folder, self.file_name)
        os.makedirs(self.folder, exist_ok=True)

        self.rectangles = []
        try:
            with open(self.data_path, "rb") as file:  # 必须使用二进制模式"rb"
                self.rectangles = pickle.load(file)
                print(f"读取数据{self.rectangles}")
        except Exception as e:
            pass

        self.root = Tk()
        self.root.geometry("400x300")
        self.button1 = Button(self.root, text=f"取新矩形({self.point_number})", command=self.take_screenshot)

        def close_window():
            self.root.destroy()  # 关闭主窗口

        self.button2 = Button(self.root, text=f"跳过({self.point_number})", command=close_window)

        self.button1.pack()
        self.button2.pack()

        self.root.mainloop()

    def take_screenshot(self):
        self.root.iconify()
        time.sleep(0.5)
        screenshot = pyautogui.screenshot()
        screenshot = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
        cv2.namedWindow("Screenshot", cv2.WINDOW_NORMAL)
        cv2.setWindowProperty("Screenshot", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        cv2.imshow("Screenshot", screenshot)
        self.draw_rectangle(screenshot)

        with open(self.data_path, "wb") as f:
            pickle.dump(self.rectangles, f)
        print(f"存储数据{self.rectangles}")

    def GetPoints(self):
        return self.rectangles

    def draw_rectangle(self, screenshot):
        start_point = None
        end_point = None
        is_dragging = False

        cnt = 0

        def mouse_callback(event, x, y, flags, param):
            nonlocal start_point, end_point, is_dragging, cnt
            self.rectangles.clear()
            if event == cv2.EVENT_LBUTTONDOWN:
                start_point = (x, y)
                is_dragging = True
            elif event == cv2.EVENT_MOUSEMOVE and is_dragging:
                end_point = (x, y)
                temp_image = screenshot.copy()
                cv2.rectangle(temp_image, start_point, end_point, (0, 255, 0), 2)
                cv2.imshow("Screenshot", temp_image)
            elif event == cv2.EVENT_LBUTTONUP:
                is_dragging = False
                end_point = (x, y)
                self.rectangles.append((start_point, end_point))
                cv2.rectangle(screenshot, start_point, end_point, (0, 255, 0), 2)
                cnt += 1
                cv2.imshow("Screenshot", screenshot)
                if cnt >= self.point_number:
                    cv2.setMouseCallback("Screenshot", lambda *args: None)
                    cv2.destroyAllWindows()
                    # self.root.deiconify()
                    self.root.destroy()

        cv2.setMouseCallback("Screenshot", mouse_callback)
        cv2.waitKey(0)


if __name__ == "__main__":
    app = ScreenshotGetPosition(2)
    print(app.GetPoints())
